/***************************************************************
*  Name:  rsa
*	    
*  Synopsis:  rsa routines
*    	    
*  Description:
*     	    
*  History:
*      	         dietrich - Jan 18, 2012:  Created.
*
*  rsa_key_t
*  rsa_init(rsa_key_t*)
*  rsa_free(rsa_key_t*)
*  rsa_decrypt(char* cipher, int size, char* decipher, rsa_key_t *)
*
*
***************************************************************/

/* $Id: rsa.c,v 1.4 2012/01/21 20:04:07 pascal Exp pascal $ */

/* Copyright (c) 2011 the authors listed at the following URL, and/or
the authors of referenced articles or incorporated external code:
http://en.literateprograms.org/Miller-Rabin_primality_test_(C,_GMP)?action=history&offset=20060620215321

Retrieved from: http://en.literateprograms.org/Miller-Rabin_primality_test_(C,_GMP)?oldid=6350
*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <gmp.h>
#include <time.h>
#include "rsa.h"
#include "utils.h"



#define COMPOSITE        0
#define PROBABLE_PRIME   1


/* -----------------------------------------------------------------
 * Helpers
 * ----------------------------------------------------------------- */

void mpz_printf_ex(int base,char * label, mpz_t bnum)
{
	printf("%s (%d bits):",label,mpz_sizeinbase(bnum,2));
	mpz_out_str(stdout,base, bnum);
	puts("");
}


void mpz_printf(char * label, mpz_t bnum)
{
	mpz_printf_ex(10,label,bnum);
	       
//	printf("%s :",label);
//	mpz_out_str(stdout,10, bnum);
//	puts("");
}

void mpz_printf2(char * label, mpz_t bnum)
{
	mpz_out_str(stdout,2, bnum);
}	


void mpz_printf16(char * label, mpz_t bnum)
{
        gmp_printf("%s (%d bits): %'02Zx\n",label,mpz_sizeinbase(bnum,2),bnum);
}





/* -----------------------------------------------------------------
 * Miller Rabin
 * ----------------------------------------------------------------- */



int miller_rabin_pass(mpz_t a, mpz_t n) {
    int i, s, result;
    mpz_t a_to_power, d, n_minus_one;

    mpz_init(n_minus_one);
    mpz_sub_ui(n_minus_one, n, 1);

    s = 0;
    mpz_init_set(d, n_minus_one);
    while (mpz_even_p(d)) {
        mpz_fdiv_q_2exp(d, d, 1);
        s++;
    }


    mpz_init(a_to_power);
    mpz_powm(a_to_power, a, d, n);
    if (mpz_cmp_ui(a_to_power, 1) == 0)  {
        result=PROBABLE_PRIME; goto exit;
    }
    for(i=0; i < s-1; i++) {
        if (mpz_cmp(a_to_power, n_minus_one) == 0) {
            result=PROBABLE_PRIME; goto exit;
        }
        mpz_powm_ui (a_to_power, a_to_power, 2, n);
    }
    if (mpz_cmp(a_to_power, n_minus_one) == 0) {
        result=PROBABLE_PRIME; goto exit;
    }
    result = COMPOSITE;

exit:
    mpz_clear(a_to_power);
    mpz_clear(d);
    mpz_clear(n_minus_one);
    return result;
}


int miller_rabin(mpz_t n, gmp_randstate_t rand_state) {
    mpz_t a;
    int repeat;
    mpz_init(a);
    for(repeat=0; repeat<20; repeat++) {
        do {
            mpz_urandomm(a, rand_state, n);
        } while (mpz_sgn(a) == 0);
        if (miller_rabin_pass(a, n) == COMPOSITE) {
            return COMPOSITE;
        }
    }
    return PROBABLE_PRIME;
}


/* -----------------------------------------------------------------
 * Generate Prime
 * ----------------------------------------------------------------- */

// generate a prime number between minbits and maxbits
int genprime (mpz_t p,  int minbits, int maxbits)
{
	mpz_t max, two;
//	mpz_t min;

	gmp_randstate_t rand_state;
	gmp_randinit_default(rand_state);
	gmp_randseed_ui (rand_state, time(NULL)+rand());
	
        mpz_init(max);
//	mpz_init(min);
        mpz_init_set_ui(two, 2);
	mpz_mul_2exp(max, two, maxbits-1);	
//	mpz_mul_2exp(min, two, minbits-1);

//	mpz_printf("MAX",max);
//	mpz_printf("MIN",min);
	
	mpz_init(p);
        do {
            mpz_urandomm(p, rand_state, max);
            if (mpz_even_p(p)) continue;
	    if (mpz_fdiv_ui(p, 3) == 0) continue;
	    if (mpz_fdiv_ui(p, 5) == 0) continue;
	    if (mpz_fdiv_ui(p, 7) == 0) continue;
	    if (mpz_fdiv_ui(p, 11) == 0) continue;
	    if (mpz_fdiv_ui(p, 13) == 0) continue;
	    if (mpz_fdiv_ui(p, 17) == 0) continue;	    	    	    
        } while ( (mpz_sizeinbase(p,2)<minbits)
		  || (miller_rabin(p, rand_state) == COMPOSITE));


	mpz_printf("PRIME",p);
	return 0;
}


/* -----------------------------------------------------------------
 * Modular inversion
 * ----------------------------------------------------------------- */

// Modular inversion.  de%n == 1
// uses the extended euclid algo.
// result : r
int mod_inv(mpz_t r,  mpz_t z, mpz_t modulus )
{
	mpz_t i, j, y1, y2, y, quotient, remainder;

	mpz_init(i);
	mpz_init(j);
	mpz_init(y1);
	mpz_init(y2);
	mpz_init(y);
	mpz_init(quotient);
	mpz_init(remainder);

	mpz_set(i,modulus);  // modulus
	mpz_set(j,z);
	mpz_set_ui(y1,1);	
	mpz_set_ui(y2,0);

	while (mpz_cmp_ui(j,0)>0) {

//		mpz_printf("i",i);
//		mpz_printf("j",j);

		mpz_fdiv_qr(quotient,remainder,i,j);
//		mpz_printf("quotient",quotient);
//		mpz_printf("remainder",remainder);
		mpz_t t;
		mpz_init(t);
		mpz_mul (t, y1, quotient);
		mpz_sub (y, y2, t);

		mpz_set(i,j);
		mpz_set(j,remainder);
		mpz_set(y2,y1);
		mpz_set(y1,y);
	}

	if (mpz_cmp_ui(i,1) )
		return -1;
	
	mpz_fdiv_r(r, y2, modulus);

	return 0;
}


/* -----------------------------------------------------------------
 * Generate RSA keys.
 * ----------------------------------------------------------------- */


int gen_rsa_keys_ex(int keylen, mpz_t exp, mpz_t modulus, mpz_t priv,mpz_t p, mpz_t q)
{
	int  pl, ql;
	int  res;

	pl = keylen /2;
	ql = keylen - pl;

	mpz_init(p);
	mpz_init(q);
	mpz_init(modulus);
	mpz_init(exp);
	mpz_init(priv);

	genprime(p, pl, pl);
	genprime(q, ql, ql);
	mpz_mul(modulus, p, q);
	mpz_set_ui(exp, 65537);

	mpz_printf("p",p);	
	mpz_printf("q",q);
	mpz_printf("modulus",modulus);
	mpz_printf("exp",exp);

	//  n=(p-1)(n-1)
	mpz_t a,b;
	mpz_init(a);
	mpz_init(b);	
	mpz_sub_ui(a,p,1);
	mpz_sub_ui(b,q,1);
	mpz_mul(a,b,a);
	res = mod_inv(priv, exp, a);
	mpz_printf("priv",priv);
	return res;
}

int gen_rsa_keys(int keylen, mpz_t exp, mpz_t modulus, mpz_t priv)
{
	mpz_t p, q;
	int ret;
	
	ret =  gen_rsa_keys_ex(keylen, exp, modulus, priv, p ,q);

	// This block is for testing only.
	{
		mpz_t tmp, dem, mod;
		mpz_init (dem);	
		mpz_init (tmp);
		mpz_init (mod);

		mpz_sub_ui(p,p,1);
		mpz_sub_ui(q,q,1);
		mpz_mul(mod,p,q);
		mpz_mul(tmp,priv,exp);
		mpz_fdiv_r(dem,tmp,mod);
		mpz_printf("Tesing d.e % modulus (should be 1)",dem);
		
	}

	return ret;
}


/* -----------------------------------------------------------------
 * rsa encrypt and decrypt
 * ----------------------------------------------------------------- */

#ifdef F_OLD_KEY
typedef struct  {
  int explen;  // size of exponent
  int modlen;  // size of modulus
  unsigned char key[1];  // key: exponent of length explen followed by modlen modulus bits
} rsa_key;
#endif

// create an mpz from an input buffer: read the buffer in the  mpz number.
void mpz_import_buf(char * buf, size_t sz, mpz_t bn)
{
	mpz_init(bn);
	mpz_import(bn, /* count*/sz, /*order*/1, /*size */1, 1, 0, buf);
}

int mpz_export_buf(char * buf, size_t sz, mpz_t bn)
{
	int bits =  mpz_sizeinbase(bn,2);

	bits = (bits+7)/8;

	// in case of null buffer return the number of bytes required
	if (buf == NULL)
		return bits;

	if (sz < bits) {
		printf("mpz_export_buff: requires %d bytes, got %d\n",bits,sz);
		return -1;
	}
	
	size_t r = 0;
	mpz_export(buf, &r, 1, 1, 1, 0, bn);

	return r;
}

#ifdef F_OLD_KEY
// import an rsa_key as exp and mod
void mpz_import_key(rsa_key * rsa, mpz_t exp, mpz_t mod)
{
	mpz_init(exp);
	mpz_init(mod);
	
	mpz_import(exp, /* count*/rsa->explen, /*order*/1, /*size */1, 1, 0, (void*)&rsa->key[0]);
	mpz_printf16("import-exp",exp);
	mpz_import(mod, /* count*/rsa->modlen, /* order*/1, /*size*/1, 1, 0, (void*)&rsa->key[rsa->explen]);
	mpz_printf16("import-mod",mod);	
}

// export mpz numbers as an rsa_key
rsa_key *  mpz_export_key(mpz_t exp, mpz_t mod)
{
	int bitse, bitsm;
	rsa_key * key;

	bitse = mpz_sizeinbase(exp,2);
	bitsm = mpz_sizeinbase(mod,2);	

	printf("bits: exp %d, mod: %d\n",bitse,bitsm);

	bitse = (bitse+7)/8;
	bitsm = (bitsm+7)/8;

	printf("bytes: exp %d, mod: %d\n",bitse,bitsm);

	key = (rsa_key *) malloc(sizeof(rsa_key)-1+bitsm+bitse);
	key->explen=bitse;
	key->modlen=bitsm;
	char * ep = key->key;
	char * em = key->key+bitse;
	size_t sz;
	mpz_export(ep, &sz, 1, 1, 1, 0, exp);
	mpz_export(em, &sz, 1, 1, 1, 0, mod);	

	return key;
	
}

rsa_key *  mpz_export_key(mpz_t exp, mpz_t mod);
#endif

int rsa_generate_key(int keylen,  rsa_key_t * public_key,  rsa_key_t * private_key)
{
	mpz_t exp, mod, priv;
	int res;

	res = gen_rsa_keys(keylen, exp, mod, priv);
	printf("rsa_generate_key: gen_rsa_key gen_rsa_keys()->%d\n",res);

	rsa_init(public_key);
	rsa_init(private_key);
	
#ifdef F_OLD_KEY
	*public_key = mpz_export_key(exp, mod);
	*private_key = mpz_export_key(priv,mod);
#endif
	mpz_set(public_key->modulus, mod);
	mpz_set(public_key->exponent, exp);

	mpz_set(private_key->modulus, mod);
	mpz_set(private_key->exponent, priv);


	// TODO : free unsused mpz!

	return 0;
}



void rsa_init(rsa_key_t * pkey)
{
  mpz_init(pkey->modulus);
  mpz_init(pkey->exponent);  
}



/* -----------------------------------------------------------------
 * RSA encryption and decryption
 * ----------------------------------------------------------------- */

int rsa_encrypt(unsigned char * input,
		unsigned int len,
		unsigned char  ** output,
		rsa_key_t * public_key)
{
	int block_size;

//	mpz_import_key(public_key, public_key->exponent, public_key->modulus);
	
	int modulus_length = mpz_sizeinbase(public_key->modulus,2);
	if (modulus_length % 8)  {
		printf("modulus size (%d) is not a multpile of 8 bits, this is not supportted\n",modulus_length);
	}
	modulus_length = (modulus_length+7)/8;
	printf("rsa_encrypt: modulus_length = %d\n",modulus_length);
	int encrypted_size = 0;
	*output=NULL;
	unsigned char * padded_block = (unsigned char *)malloc(modulus_length);
	while (len)
	{
	  printf("Block %d\n",encrypted_size/modulus_length);
		encrypted_size += modulus_length;
		block_size  = (len < modulus_length -11) ? len : (modulus_length -11);
		printf("rsa_encrypt: len %d, encrypted_size %d, block_size %d\n",len,encrypted_size,block_size);
		memset (padded_block, 0, modulus_length);
		memcpy (padded_block + (modulus_length - block_size), input, block_size);
		padded_block[1] = 0x02;
		int i;
		for (i = 2; i< (modulus_length - block_size -1); i++)
			padded_block[i] = 0x03; //rand();
		memdump("padded-input",padded_block,modulus_length);
		mpz_t m;
		mpz_t rop;
		mpz_init(rop);
		mpz_import_buf(padded_block, modulus_length, m);
		mpz_printf16("encrypted-text", m);		
		mpz_powm (rop, m, public_key->exponent, public_key->modulus);
		mpz_printf16("encrypted-block", rop);
		*output = (unsigned char *)realloc(*output, encrypted_size);
		int r = mpz_export_buf(*output+(encrypted_size - modulus_length), modulus_length, rop);
		int sh = modulus_length - r;
		if ( sh > 0) {
		  unsigned char *p = *output+(encrypted_size - modulus_length);
		  printf("shifting right %d bytes\n",r);
		  memmove(p+sh,p,r);
		  memset(p,0,sh);
		}
		printf("mpz_export_buf: ret=%d\n",r);
		memdump("encrypted text",*output+encrypted_size-modulus_length,modulus_length);
		len -= block_size;
		input += block_size;
		mpz_clear(m);
		mpz_clear(rop);
	}
	printf("finished, encrypted_size %d\n",encrypted_size);
	free (padded_block);
	return encrypted_size;
}



int rsa_decrypt(unsigned char *input,
		unsigned int len,
		unsigned char ** output,
		rsa_key_t * private_key)
{
	int out_len=0;

	mpz_printf16("priv-exp",private_key->exponent);
	mpz_printf16("priv-mod",private_key->modulus);
	int modulus_length = mpz_sizeinbase(private_key->modulus,2);
	if (modulus_length % 8)  {
		printf("modulus size (%d) is not a multpile of 8 bits, this is not supportted\n",modulus_length);
	}
	modulus_length = (modulus_length+7)/8;	
	if (len % modulus_length) {
		printf("rsa_decrypt:  input must be multiple of modulus len %d/%d\n",len,modulus_length);
	}
	printf("rsa_decrypt: modulus_length = %d\n",modulus_length);
	unsigned char * padded_block = (unsigned char *)malloc(modulus_length);
	*output = NULL;
	int block=0;
	while (len) {
	        printf("iteration %d\n",block++);
		if (len < modulus_length)
			return -1;
		mpz_t c, rop;
		mpz_init(rop);
		mpz_import_buf(input, modulus_length, c);
		mpz_printf16("cipher-text", c);		
		mpz_powm (rop, c, private_key->exponent, private_key->modulus);
		mpz_printf16("decrypted", rop);
		int l=mpz_sizeinbase(rop,2);
		l=(l+7)/8;
		printf("decrypted size = %d\n",l);
		mpz_export_buf(padded_block+(modulus_length-l), l /*modulus_length*/, rop);

		if (padded_block[1] > 2) {
			printf("rsa_decrypt: decryption error or unrecognized block type %d.\n",padded_block[1]);
			memdump("padded_block",padded_block,modulus_length);
			mpz_clear(rop);
			mpz_clear(c);
			free(padded_block);
			return -1;
		}
		// Find first 0 byte after the padding type byte, this signifies start of data.
		int i=2;
		while (padded_block[i++] && i<modulus_length);
		if (i == modulus_length) {
			printf("rsa_decrypt: end of padding not found\n");
			return -1;
		}
		out_len += modulus_length -i;
		*output = (unsigned char *)realloc(*output, out_len);
		memcpy(*output+out_len-modulus_length+i,padded_block+i, modulus_length - i);
		len -= modulus_length;
		input += modulus_length;
		mpz_clear(c);
		mpz_clear(rop);
	}
	free(padded_block);
	return out_len;
}


